 #!c:\\python\\python.exe

# smashsniffer v01, instantly sniff and log mashed potatos
# - James Kearney / kitkat
#
# monitor a windows process / service with threads, automatically restart if it suicides or crashes and log 
# - crashbins to a crashbin file
# - program stops and restart it when it dies
# - addresses and stack unwind of breakpoint, access_violation, new_thread
# - (optionally) log the last X seconds of packets (bpf filtered) to a file corresponding to crash
#
#   written by kitkat ripped from parts of pydbg, sulley, taof, crashbin_explorer, 
#   rod's head (cheers rod), and mark's head / help debugging, 
#   ps s/$/%/g - jjk stop writing perl all the time
#
# v01   -   27/01/2011
# v02   -   04/02/2011
#       -   now does a packet dump of [window] seconds on process suicide as well
import os
import sys
import time
import utils
import getopt
import datetime
import traceback
import threading
import subprocess

sys.path.append(r"..\..\..\paimei")
sys.path.append(r"/kitkat/blah/directory/paimei-read-only/")

import pydbg
import pcapy
import impacket

from pydbg         import *
from pydbg.defines import *
from pcapy import findalldevs, open_live
from datetime import datetime, timedelta
from impacket import ImpactDecoder, ImpactPacket
from impacket.ImpactDecoder import EthDecoder, LinuxSLLDecoder
# import debug

############################################################
## decoder thread, decodes and records packets
class DecoderThread(threading.Thread):
    def __init__(self, pack_mon ):
        self.decoder = None
        self.stack = {}
        self.pack_mon = pack_mon
        self.pcap = self.pack_mon.pcapobj
        self.seconds = self.pack_mon.seconds

        # Query the type of the link and instantiate a decoder accordingly.
        datalink = self.pcap.datalink()
        if pcapy.DLT_EN10MB == datalink:
            self.decoder = EthDecoder()
        elif pcapy.DLT_LINUX_SLL == datalink:
            self.decoder = LinuxSLLDecoder()
        else:
            raise Exception("Datalink type not supported: " % datalink)
        threading.Thread.__init__(self)
        print "DecoderThread initialised\n"

    def run(self):
        print "decoderthread run\n"
        self.pcap.loop(0, self.packetHandler)
    
    def getStack( self ):
        return self.stack

    def writepcap( self ):
        try:
            now = datetime.now()
            if not len( self.stack.keys() ) == 0:
                fn = "capture.crashdumped.%s-%s-%s_%s-%s-%s.pcap" % ( str(now.year), str(now.month), str(now.day), str(now.hour), str(now.minute), str(now.second) )
                dumper = self.pcap.dump_open(fn)
                for thetime in self.stack.keys():
                    ( thehdr, thedata ) = self.stack[thetime]
                    dumper.dump( thehdr, thedata )
        except Exception, e:
            print "exception: %s" % str( e )
            traceback.print_exc( file=sys.stdout )
    
    def packetHandler(self, hdr, data):
        # stack of packets for the last X seconds
        now = datetime.now() 
        self.stack[now] = ( hdr, data )
        for key in self.stack.keys():
            dthen = key
            if ( now - dthen ) > timedelta( seconds = int(self.seconds) ):
                del self.stack[dthen]

        ether = self.decoder.decode(data)
        iphdr = ether.child()
        tcphdr = iphdr.child()
        # print "ether:\n" + str(ether)
        # print "iphdr:\n" + str(iphdr)
        # print "tcphdr:\n" + str(tcphdr)
        srcip = iphdr.get_ip_src()
        dstip = iphdr.get_ip_dst()
        print "\nconnection from %s to %s\n" % ( str(srcip), str(dstip) )
        print "\n %s" % str( self.decoder.decode(data) )


############################################################
## maintains a stack of packets for a set timeperiod
class pack_mon():
    def __init__( self, procmon ):
        self.proc_mon = procmon
        self.interface = int(self.proc_mon.interface)
        self.bpf = self.proc_mon.bpf
        self.seconds = self.proc_mon.seconds

        ifs = findalldevs()
        self.dev = ifs[self.interface]

        print "\n[pack_mon] monitoring interface %s with filter %s\n" % ( self.interface, self.bpf )
        self.pcapobj = open_live(self.dev, 1500, 0, 100)
        self.pcapobj.setfilter(self.bpf)
        print "Listening on %s: net=%s, mask=%s, linktype=%d" % (self.dev, self.pcapobj.getnet(), self.pcapobj.getmask(), self.pcapobj.datalink())
        # Start sniffing thread and finish main thread.
        self.decoder = DecoderThread( self )
        self.decoder.setDaemon( True )
        self.decoder.start()
        print "done pack_mon init\n"

    def writepcap( self ):
        if self.decoder:
            self.decoder.writepcap()

#######################################################
## Does the process monitoring / crash dumping / restarting
## 
class proc_mon( threading.Thread ):
    def __init__( self, proc_name, start_cmd, crash_filename, log_filename, interface, bpf, seconds ):
        threading.Thread.__init__(self)
        self.proc_name = proc_name.lower()
        self.debugger = pydbg()
        self.attached = False
        self.pid = -1
        self.start_cmd = start_cmd
        self.crash_bin = utils.crash_binning.crash_binning()
        self.crash_file = crash_filename
        self.log_file = log_filename

        self.pack_mon = None
        self.interface = interface
        self.bpf = bpf
        self.seconds = seconds

        self.setName( "%d" % time.time() )

        self.print_file( "[proc_mon][BEGIN] new monitor thread starting\n\n" )

        if not self.interface == -1 and self.bpf:
            self.pack_mon = pack_mon( self )
            # self.pack_mon.setDaemon( True )
            # self.pack_mon.start()

        if not os.access(os.path.dirname(self.crash_file), os.X_OK):
            grumble = "[!fatal][access crash_file] invalid path specified for crash bin (give filename): %s\n" % self.crash_file
            self.print_file( grumble )
            #raise Exception
        if not os.access(os.path.dirname(self.log_file), os.X_OK):
            grumble = "[!fatal][access log_file] invalid path specified for log file (give filename): %s\n" % self.log_file
            self.print_file( grumble )
            #raise Exception

        try:
            if os.path.isfile( self.crash_file): 
                self.crash_bin.import_file( self.crash_file )
                if not self.crash_bin or not self.crash_bin.bins:
                    self.print_file( "\n[inf] crash_bin len: 0\n" )
                else:
                    bincount = len(self.crash_bin.bins)
                    count = bincount
                    if not bincount or bincount==0:
                        count = 0
                    grumble = "[inf] crash_bin len: %s\n" % str(count)
                    self.print_file( grumble )
                    if count != 0:
                        synopsis="no synopsis set"
                        for bin, crashes in self.crash_bin.bins.iteritems():
                            synopsis = self.crash_bin.crash_synopsis( crashes[0] ).split("\n")[0]
                        grumble = "[inf] crash_synopsis: *** %s ***\n" % synopsis
                        self.print_file( grumble )
            else:
                self.print_file( "[import crashbin] - not importing, doesn't exist\n" )
        except Exception, e:
            print "except: %s" % str(e)
            self.print_file( "[!fatal][import crash bin] %s\n" % str(e) )
            #traceback.print_exc( file=sys.stdout )
            pass

        # self.debugger.set_callback( EXCEPTION_BREAKPOINT, debug.handler_breakpoint )
        self.debugger.set_callback( EXCEPTION_ACCESS_VIOLATION, self.handler_access_violation )
        self.debugger.set_callback( CREATE_THREAD_DEBUG_EVENT, self.handler_new_thread )
        # self.debugger.set_callback( USER_CALLBACK_DEBUG_EVENT, self.handler_debug_event )


    def handler_access_violation ( self, dbg ):
        #if pydbg.dbg.u.Exception.dwFirstChance:
        #    return DBG_EXCEPTION_NOT_HANDLED
        outstring = "[crash][handler_access_violation]\n" 
        outstring = outstring + "[dwThreadId] called from thread %d \n" % dbg.dbg.dwThreadId
        outstring = outstring + "[exception_address] %08x \n" % dbg.exception_address
        outstring = outstring + "[dump_context]\n" + dbg.process_snapshot().dump_context() + "\n"
        outstring = outstring + "[endcrash]\n\n"
        self.print_file( outstring )

        self.crash_bin.record_crash( dbg )
        self.crash_bin.export_file( self.crash_file )
    
        if self.pack_mon:
            self.pack_mon.writepcap()

        dbg.terminate_process()
        # return DBG_CONTINUE
        return DBG_EXCEPTION_NOT_HANDLED

    def handler_new_thread ( self, dbg ):
        outstring = "[newthread][handler_new_thread]\n" 
        outstring = outstring + "[dwThreadId] called from thread %d \n" % dbg.dbg.dwThreadId
        outstring = outstring + "[endnewthread]\n\n"
        self.print_file( outstring )
        return DBG_CONTINUE

    def handler_breakpoint ( self, dbg ):
        if dbg.first_breakpoint:
            return DBG_CONTINUE
        oustring = "[breakpoint][handler_breakpoint]\n" 
        outstring = outstring + "[dwThreadId] called from thread %d \n" % dbg.dbg.dwThreadId
        outstring = outstring + "[exception_address] %08x \n" % dbg.execption_address
        outstring = outstring + "[dump_context]\n" + self.process_snapshot().dump_context() + "\n"
        outstring = outstring + "[endbreakpoint]\n\n" 
        self.print_file( outstring )
        return DBG_EXCEPTION_NOT_HANDLED


    def handler_debug_event (self, *args):
        print "debug event\n"

    def print_file( self, stringer ):
        try: 
            now = datetime.now() 
            FILE = open( self.log_file, "a" )
            output = "[ %s ] : %s" % (str(now), stringer)
            FILE.writelines( output )
            FILE.close()
            print "%s" % output
        except Exception, e:
            # self.print_file( "except: %s" % str( e ) )
            print "string was %s" % stringer
            traceback.print_exc( file=sys.stdout )
            sys.exit( 2 )

    def run( self ): 
        self.print_file( "[threadstart] debugger thread-%s looking for process name: %s\n" % (self.getName(), self.proc_name) )
        # watch for and try attaching to the process.
        try:
            self.watch()
            self.debugger.attach(self.pid)
            self.debugger.run()
            self.print_file( "[threadend] debugger thread-%s exiting\n" % self.getName() )
#            if self.pack_mon:   
#                self.pack_mon
        except Exception, e:
            print "except: %s" % str(e)
            pass
        # XXX - removing the following line appears to cause some concurrency issues.
        del(self.debugger)

    
    def watch( self ):
        # block on don't know pid
        while self.pid==-1:
            for (pid, name) in self.debugger.enumerate_processes():
                if name.lower() == self.proc_name.lower():
                    self.pid = pid
                    break
            if self.pid == -1:
                self.attached = False
                self.print_file( "[procstart] didn't find %s, so starting \"%s\"\n" % ( self.proc_name, self.start_cmd ) )
                subprocess.Popen( start_cmd )
                time.sleep( 2 )
        self.print_file( "[procfound] debugger thread-%s found match on pid %d" % (self.getName(), self.pid) ) 


ERR   = lambda msg: sys.stderr.write("ERR> " + msg + "\n") or sys.exit(1)
USAGE = "USAGE: process_monitor.py"                                                                \
        "\n    [-p|--proc_name\t\"str\"]\tprocess name to search for "                          \
        "\n    [-s|--start_cmd\t\"str\"]\tstart commands to execute for starting process"       \
        "\n    [-c|--crash_file\t\"str\"]\tcrashbin file to log to (crashbin format)"            \
        "\n    [-l|--log_file\t\"str\"]\tfilename to log to (text logging format)"             \
        "\n    [-i|--interface\t\"int\"]\tinterface number to capture packets on (optional)"    \
        "\n    [-f|--filter\t\"str\"]\tthe bpf filter for packet capture (tcpdump style)"    \
        "\n    [-w|--window\t\"int\"]\tkeep stack of packets <int> no. of secs (to create pcap when crash) [default 10 secs]"     \
        "\n    [-d|--dump\t\t]\tshow available interfaces to capture on\n"            \
        "\n    - version .02 kitkat 04/02/2011\n"

def show_interfaces():
    ifs = findalldevs()
    if len(ifs)==0:
        print "cannot see any devs - got permissions?\n"
        sys.exit( 2 )
    count=0
    for iface in ifs:
        print '%i - %s' % ( count, ifs )
        count+=1
        

if __name__ == "__main__":
    proc_name=None
    start_cmd=None
    crash_file=None
    log_file=None
    proc_re=None
    interface=-1
    bpf=None
    dump=False
    seconds=10
    try:
        opts, args = getopt.getopt(sys.argv[1:], "p:s:c:l:i:f:dw:", ["proc_name=", "start_cmd=","crash_file=","log_file=","interface=-1","filter=","dump","window="])
    except getopt.GetoptError:
        ERR(USAGE)
    for opt, arg in opts:
        if opt in ("-p", "--proc_name"):   proc_name  = arg
        if opt in ("-s", "--start_cmd"):   start_cmd  = arg
        if opt in ("-c", "--crash_file"):  crash_file = arg
        if opt in ("-l", "--log_file"):    log_file   = arg
        if opt in ("-i", "--interface"):   interface  = arg
        if opt in ("-f", "--filter"):      bpf  = arg
        if opt in ("-w", "--window"):      seconds  = arg
        if opt in ("-d", "--dump"):        dump = True

    if dump:
        show_interfaces()
        sys.exit( 1 )
    if not proc_name or not start_cmd or not crash_file or not log_file:
        print "available interfaces:\n"
        show_interfaces()
        print "\n"
        ERR(USAGE)
    if interface and not bpf:
        print "if you define interface, also define a capture filter\n"
        ERR(USAGE)

    try:
        while True: 
            if proc_re and proc_re.isAlive():
                while proc_re.isAlive():
                    # print "thread still alive"
                    time.sleep(1)
            # print "new thread initialise"
            proc_re = proc_mon( proc_name, start_cmd, crash_file, log_file, interface, bpf, seconds )
            proc_re.setDaemon( True )
            proc_re.start()
            time.sleep(1)
    except Exception as e:
        print "except: %s" % str(e)
        traceback.print_exc( file=sys.stdout )


